import torch
import torch.nn as nn
from .yolov1_basic import Conv


# Spatial Pyramid Pooling
class SPPF(nn.Module):
    """
        该代码参考YOLOv5的官方代码实现 https://github.com/ultralytics/yolov5
    """
    def __init__(self, in_dim, out_dim, expand_ratio=0.5, pooling_size=5, act_type='lrelu', norm_type='BN'):
        super().__init__()
        inter_dim = int(in_dim * expand_ratio)
        self.out_dim = out_dim
        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
        self.cv2 = Conv(inter_dim * 4, out_dim, k=1, act_type=act_type, norm_type=norm_type)
        self.m = nn.MaxPool2d(kernel_size=pooling_size, stride=1, padding=pooling_size // 2)

    def forward(self, x):
        x = self.cv1(x) # 1*1卷积降维 [B, C, H, W] -> [B, C/2, H, W]
        y1 = self.m(x) # 第一次池化 [B, C/2, H, W]
        y2 = self.m(y1) # 第二次池化 [B, C/2, H, W]
        y3 = self.m(y2) # 第三次池化 [B, C/2, H, W]

        # torch.cat((x, y1, y2, y3), 1) 将四个特征图在通道维度上拼接
        # 即在通道方向上append起来，通道数变多了，但是特征图形状不变
        # 结果形状 [B, C/2*4, H, W]

        return self.cv2(torch.cat((x, y1, y2, y3), 1))
    
# SPPF多级池化为什么能增加感受野
    # 池化操作的本质：
        # 虽然输入输出大小相同，但每次池化都是在更大的感受野上进行操作
        # 第二次池化是在第一次池化的结果上进行的，所以它"看到"的范围更大
    # 具体例子，使用5x5的池化核（kernel_size=5）
        # 第一次池化 y1 = self.m(x):
        # - 在原始特征图上，每个输出点是从5x5的局部区域取最大值
        # - 感受野：5x5

        # 第二次池化 y2 = self.m(y1):
        # - 在y1上进行池化，每个输出点仍然是从5x5的区域取最大值
        # - 但是y1的每个点已经包含了原始图像5x5区域的信息
        # - 所以y2的每个点实际上覆盖了原始图像9x9的区域
        # - 感受野：9x9

        # 第三次池化 y3 = self.m(y2):
        # - 在y2上进行池化
        # - 每个点覆盖原始图像13x13的区域
        # - 感受野：13x13
    # 为什么有效：
        # 虽然空间分辨率相同，但每次池化都让网络能够捕获更大范围的特征
        # 这有助于检测不同尺度的目标：
            # 第一次池化：适合检测小目标
            # 第二次池化：适合检测中等大小的目标
            # 第三次池化：适合检测大目标
    # 实际效果：
        # 原始特征图：关注局部细节
        # 第一次池化后：可以看到中等范围的特征
        # 第二次池化后：可以看到更大范围的特征
        # 第三次池化后：可以看到最大范围的特征
    # 所以，虽然空间大小不变，但每次池化都让网络能够"看到"更大范围的图像信息，这对于目标检测来说是非常有价值的。


# 搭建Neck网络
def build_neck(cfg, in_dim, out_dim):
    model = cfg['neck']
    print('==============================')
    print('Neck: {}'.format(model))
    # build neck
    if model == 'sppf':
        neck = SPPF(
            in_dim=in_dim,
            out_dim=out_dim,
            expand_ratio=cfg['expand_ratio'], 
            pooling_size=cfg['pooling_size'],
            act_type=cfg['neck_act'],
            norm_type=cfg['neck_norm']
            )
    else:
        raise NotImplementedError('Neck {} not implemented.'.format(cfg['neck']))

    return neck
        